#!/usr/bin/env pwsh
<#
.SYNOPSIS
    Enhanced Training API测试脚本 - 测试后端增强训练API
.DESCRIPTION
    测试LightGBM + XGBoost + 智能集成训练API的完整功能
.AUTHOR
    Claude AI Assistant
.DATE
    2025-07-19
#>

# 配置API基础URL
$BaseURL = "http://localhost:8000/api/v1/enhanced-training"
$Headers = @{
    "Content-Type" = "application/json"
    "Accept" = "application/json"
}

# 颜色输出函数
function Write-ColorOutput {
    param(
        [string]$Message,
        [string]$Color = "White"
    )
    Write-Host $Message -ForegroundColor $Color
}

function Write-Success {
    param([string]$Message)
    Write-ColorOutput "✅ $Message" "Green"
}

function Write-Error {
    param([string]$Message)
    Write-ColorOutput "❌ $Message" "Red"
}

function Write-Info {
    param([string]$Message)
    Write-ColorOutput "ℹ️  $Message" "Cyan"
}

function Write-Warning {
    param([string]$Message)
    Write-ColorOutput "⚠️  $Message" "Yellow"
}

# 测试API连接
function Test-APIConnection {
    Write-Info "测试API连接..."
    try {
        $response = Invoke-RestMethod -Uri "$BaseURL/models/status" -Method GET -Headers $Headers -TimeoutSec 10
        Write-Success "API连接成功"
        Write-Info "Available models: $($response.models.Count)"
        return $true
    }
    catch {
        Write-Error "API连接失败: $($_.Exception.Message)"
        return $false
    }
}

# 获取模型状态
function Get-ModelsStatus {
    Write-Info "获取模型状态..."
    try {
        $response = Invoke-RestMethod -Uri "$BaseURL/models/status" -Method GET -Headers $Headers
        if ($response.success) {
            Write-Success "模型状态获取成功"
            foreach ($model in $response.models) {
                Write-Host "  - $($model.name) ($($model.key)): $($model.status)" -ForegroundColor Yellow
                Write-Host "    $($model.description)" -ForegroundColor Gray
            }
        }
        return $response
    }
    catch {
        Write-Error "获取模型状态失败: $($_.Exception.Message)"
        return $null
    }
}

# 启动训练
function Start-Training {
    param(
        [string]$TrainRange = "3_months",
        [bool]$EnableParallel = $true,
        [bool]$EnableShap = $true
    )
    
    Write-Info "启动增强训练 (数据范围: $TrainRange, 并行: $EnableParallel, SHAP: $EnableShap)..."
    
    $trainingConfig = @{
        train_range = $TrainRange
        enable_parallel = $EnableParallel
        enable_shap = $EnableShap
        dynamic_weights = $true
        models = @("lightgbm", "xgboost", "ensemble")
        lightgbm = @{
            learning_rate = 0.05
            num_leaves = 31
            feature_fraction = 0.9
            bagging_fraction = 0.8
            lambda_l1 = 0.1
            lambda_l2 = 0.1
        }
        xgboost = @{
            learning_rate = 0.05
            max_depth = 4
            n_estimators = 100
            subsample = 0.6
            colsample_bytree = 0.6
            reg_alpha = 0.5
            reg_lambda = 0.5
        }
    }
    
    $jsonBody = $trainingConfig | ConvertTo-Json -Depth 10
    
    try {
        $response = Invoke-RestMethod -Uri "$BaseURL/train/start" -Method POST -Headers $Headers -Body $jsonBody -TimeoutSec 30
        if ($response.success) {
            Write-Success "训练任务启动成功"
            Write-Info "预计耗时: $($response.estimated_time)"
            return $true
        } else {
            Write-Error "训练启动失败: $($response.error)"
            return $false
        }
    }
    catch {
        Write-Error "启动训练失败: $($_.Exception.Message)"
        if ($_.Exception.Response) {
            $errorDetails = $_.Exception.Response.Content.ReadAsStringAsync().Result
            Write-Error "错误详情: $errorDetails"
        }
        return $false
    }
}

# 监控训练状态
function Monitor-TrainingStatus {
    param(
        [int]$MaxWaitTime = 600,  # 10分钟
        [int]$CheckInterval = 5   # 5秒检查一次
    )
    
    Write-Info "开始监控训练状态..."
    $startTime = Get-Date
    $lastProgress = -1
    
    while ((Get-Date) - $startTime -lt [TimeSpan]::FromSeconds($MaxWaitTime)) {
        try {
            $response = Invoke-RestMethod -Uri "$BaseURL/train/status" -Method GET -Headers $Headers
            
            if ($response.success) {
                $isTraining = $response.is_training
                $progress = $response.progress
                $currentModel = $response.current_model
                
                # 只在进度变化时显示
                if ($progress -ne $lastProgress) {
                    $elapsedTime = (Get-Date) - $startTime
                    Write-Host "[$($elapsedTime.ToString('mm\:ss'))] 训练进度: $progress% (当前模型: $currentModel)" -ForegroundColor Cyan
                    $lastProgress = $progress
                    
                    # 显示各模型状态
                    Write-Host "  模型状态:" -ForegroundColor Gray
                    foreach ($model in $response.models_status.PSObject.Properties) {
                        $status = switch ($model.Value) {
                            "pending" { "⏳ 等待中" }
                            "training" { "🔄 训练中" }
                            "completed" { "✅ 已完成" }
                            "failed" { "❌ 失败" }
                            default { "❓ 未知" }
                        }
                        Write-Host "    $($model.Name): $status" -ForegroundColor Gray
                    }
                }
                
                # 显示最新日志
                if ($response.logs -and $response.logs.Count -gt 0) {
                    $latestLogs = $response.logs | Select-Object -Last 3
                    foreach ($log in $latestLogs) {
                        $logColor = switch ($log.type) {
                            "success" { "Green" }
                            "error" { "Red" }
                            "warning" { "Yellow" }
                            default { "White" }
                        }
                        Write-Host "  [$($log.time)] $($log.message)" -ForegroundColor $logColor
                    }
                }
                
                if (-not $isTraining) {
                    Write-Success "训练已完成!"
                    return $response
                }
            }
        }
        catch {
            Write-Warning "获取训练状态失败: $($_.Exception.Message)"
        }
        
        Start-Sleep -Seconds $CheckInterval
    }
    
    Write-Warning "训练监控超时 ($MaxWaitTime 秒)"
    return $null
}

# 获取训练结果
function Get-TrainingResults {
    Write-Info "获取训练结果..."
    try {
        $response = Invoke-RestMethod -Uri "$BaseURL/train/results" -Method GET -Headers $Headers
        
        if ($response.success) {
            Write-Success "训练结果获取成功"
            
            Write-Host "`n📊 训练结果摘要:" -ForegroundColor Yellow
            Write-Host "  总模型数: $($response.summary.total_models)" -ForegroundColor Gray
            Write-Host "  完成模型数: $($response.summary.completed_models)" -ForegroundColor Gray
            Write-Host "  训练总耗时: $([math]::Round($response.summary.training_time, 2))秒" -ForegroundColor Gray
            
            Write-Host "`n📈 各模型性能:" -ForegroundColor Yellow
            foreach ($model in $response.results.PSObject.Properties) {
                $metrics = $model.Value
                if ($metrics.mae -and $metrics.mae -gt 0) {
                    Write-Host "  $($model.Name.ToUpper()):" -ForegroundColor Cyan
                    Write-Host "    MAE:  $([math]::Round($metrics.mae, 4))" -ForegroundColor White
                    Write-Host "    RMSE: $([math]::Round($metrics.rmse, 4))" -ForegroundColor White
                    Write-Host "    R²:   $([math]::Round($metrics.r2, 4))" -ForegroundColor White
                    Write-Host "    MAPE: $([math]::Round($metrics.mape, 2))%" -ForegroundColor White
                    
                    if ($metrics.weights) {
                        Write-Host "    权重配置:" -ForegroundColor Gray
                        foreach ($weight in $metrics.weights.PSObject.Properties) {
                            Write-Host "      $($weight.Name): $([math]::Round($weight.Value * 100, 1))%" -ForegroundColor Gray
                        }
                    }
                }
            }
            
            return $response
        } else {
            Write-Error "获取训练结果失败: $($response.error)"
            return $null
        }
    }
    catch {
        Write-Error "获取训练结果失败: $($_.Exception.Message)"
        return $null
    }
}

# 获取模型可解释性分析
function Get-ModelInterpretability {
    param([string]$ModelName)
    
    Write-Info "获取 $ModelName 模型可解释性分析..."
    try {
        $response = Invoke-RestMethod -Uri "$BaseURL/train/interpretability/$ModelName" -Method GET -Headers $Headers
        
        if ($response.success) {
            Write-Success "$ModelName 可解释性分析获取成功"
            
            $interpretability = $response.interpretability
            if ($interpretability -and $interpretability.interpretability_analysis) {
                Write-Host "`n🔍 $($ModelName.ToUpper()) 可解释性分析:" -ForegroundColor Yellow
                foreach ($analysis in $interpretability.interpretability_analysis.PSObject.Properties) {
                    Write-Host "  $($analysis.Name): $($analysis.Value)" -ForegroundColor White
                }
                
                Write-Host "`n📊 特征贡献度:" -ForegroundColor Yellow
                if ($interpretability.historical_load) {
                    Write-Host "  历史负荷: $([math]::Round($interpretability.historical_load.contribution_ratio * 100, 1))%" -ForegroundColor Cyan
                }
                if ($interpretability.weather) {
                    Write-Host "  气象因素: $([math]::Round($interpretability.weather.contribution_ratio * 100, 1))%" -ForegroundColor Cyan
                }
                if ($interpretability.time) {
                    Write-Host "  时间因素: $([math]::Round($interpretability.time.contribution_ratio * 100, 1))%" -ForegroundColor Cyan
                }
            }
            
            return $response
        } else {
            Write-Error "获取可解释性分析失败: $($response.error)"
            return $null
        }
    }
    catch {
        Write-Error "获取可解释性分析失败: $($_.Exception.Message)"
        return $null
    }
}

# 压力测试
function Start-StressTest {
    param([int]$ConcurrentRequests = 3)
    
    Write-Info "开始压力测试 (并发请求数: $ConcurrentRequests)..."
    
    $jobs = @()
    for ($i = 1; $i -le $ConcurrentRequests; $i++) {
        $job = Start-Job -ScriptBlock {
            param($BaseURL, $Headers, $RequestId)
            
            try {
                $response = Invoke-RestMethod -Uri "$BaseURL/train/status" -Method GET -Headers $Headers -TimeoutSec 10
                return @{
                    RequestId = $RequestId
                    Success = $response.success
                    ResponseTime = (Measure-Command { 
                        Invoke-RestMethod -Uri "$BaseURL/train/status" -Method GET -Headers $Headers -TimeoutSec 10 
                    }).TotalMilliseconds
                }
            }
            catch {
                return @{
                    RequestId = $RequestId
                    Success = $false
                    Error = $_.Exception.Message
                    ResponseTime = -1
                }
            }
        } -ArgumentList $BaseURL, $Headers, $i
        
        $jobs += $job
    }
    
    Write-Info "等待并发请求完成..."
    $results = $jobs | Wait-Job | Receive-Job
    $jobs | Remove-Job
    
    $successCount = ($results | Where-Object { $_.Success }).Count
    $avgResponseTime = ($results | Where-Object { $_.ResponseTime -gt 0 } | Measure-Object -Property ResponseTime -Average).Average
    
    Write-Host "`n🔥 压力测试结果:" -ForegroundColor Yellow
    Write-Host "  总请求数: $ConcurrentRequests" -ForegroundColor Gray
    Write-Host "  成功请求数: $successCount" -ForegroundColor Green
    Write-Host "  失败请求数: $($ConcurrentRequests - $successCount)" -ForegroundColor Red
    Write-Host "  平均响应时间: $([math]::Round($avgResponseTime, 2))ms" -ForegroundColor Cyan
    
    return $results
}

# 完整测试流程
function Start-CompleteTest {
    param(
        [string]$TrainRange = "1_month",
        [bool]$EnableStressTest = $false
    )
    
    Write-Host "🚀 开始增强训练API完整测试" -ForegroundColor Magenta
    Write-Host "=" * 60 -ForegroundColor Gray
    
    # 1. 测试连接
    if (-not (Test-APIConnection)) {
        Write-Error "API连接失败，测试终止"
        return
    }
    
    # 2. 获取模型状态
    $modelsStatus = Get-ModelsStatus
    if (-not $modelsStatus) {
        Write-Warning "无法获取模型状态，继续测试..."
    }
    
    # 3. 启动训练
    $trainingStarted = Start-Training -TrainRange $TrainRange -EnableParallel $true -EnableShap $true
    if (-not $trainingStarted) {
        Write-Error "训练启动失败，测试终止"
        return
    }
    
    # 4. 监控训练状态
    $finalStatus = Monitor-TrainingStatus -MaxWaitTime 600 -CheckInterval 5
    if (-not $finalStatus) {
        Write-Warning "训练监控超时，尝试获取当前结果..."
    }
    
    # 5. 获取训练结果
    $results = Get-TrainingResults
    if ($results) {
        # 6. 获取各模型可解释性分析
        foreach ($modelName in @("lightgbm", "xgboost")) {
            Get-ModelInterpretability -ModelName $modelName | Out-Null
        }
    }
    
    # 7. 压力测试（可选）
    if ($EnableStressTest) {
        Start-StressTest -ConcurrentRequests 3 | Out-Null
    }
    
    Write-Host "`n🎉 测试完成!" -ForegroundColor Magenta
    Write-Host "=" * 60 -ForegroundColor Gray
    
    return @{
        ConnectionTest = $true
        ModelsStatus = $modelsStatus
        TrainingStarted = $trainingStarted
        TrainingResults = $results
    }
}

# 主程序入口
function Main {
    param(
        [string]$Action = "complete",
        [string]$TrainRange = "1_month",
        [bool]$EnableStressTest = $false
    )
    
    switch ($Action.ToLower()) {
        "connection" { Test-APIConnection }
        "status" { Get-ModelsStatus }
        "start" { Start-Training -TrainRange $TrainRange }
        "monitor" { Monitor-TrainingStatus }
        "results" { Get-TrainingResults }
        "interpretability" { 
            Get-ModelInterpretability -ModelName "lightgbm"
            Get-ModelInterpretability -ModelName "xgboost"
        }
        "stress" { Start-StressTest }
        "complete" { Start-CompleteTest -TrainRange $TrainRange -EnableStressTest $EnableStressTest }
        default { 
            Write-Error "未知操作: $Action"
            Write-Info "可用操作: connection, status, start, monitor, results, interpretability, stress, complete"
        }
    }
}

# 脚本参数处理
param(
    [string]$Action = "complete",
    [string]$TrainRange = "1_month",
    [switch]$StressTest
)

# 执行主程序
Main -Action $Action -TrainRange $TrainRange -EnableStressTest $StressTest.IsPresent